package main

import (
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"
	"time"

	"wgmat/compute"
	"wgmat/gpucore"
	"wgmat/matrix"

	"gonum.org/v1/gonum/mat"
)

// matrixEqual 检查两个矩阵是否在指定容差内相等
func matrixEqual(a, b mat.Matrix, tol float64) bool {
	ar, ac := a.Dims()
	br, bc := b.Dims()
	if ar != br || ac != bc {
		return false
	}
	for i := 0; i < ar; i++ {
		for j := 0; j < ac; j++ {
			diff := math.Abs(a.At(i, j) - b.At(i, j))
			if diff > tol {
				return false
			}
		}
	}
	return true
}

// vecEqual 检查两个向量是否在指定容差内相等
func vecEqual(a, b []float32, tol float64) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		diff := math.Abs(float64(a[i]) - float64(b[i]))
		if diff > tol {
			return false
		}
	}
	return true
}

// gonumMatrixToSlice 将 gonum 矩阵转换为 []float32 切片
func gonumMatrixToSlice(m mat.Matrix) []float32 {
	rows, cols := m.Dims()
	data := make([]float32, rows*cols)
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			data[i*cols+j] = float32(m.At(i, j))
		}
	}
	return data
}

// sliceToGonumMatrix 将 []float32 切片转换为 gonum 矩阵
func sliceToGonumMatrix(data []float32, rows, cols int) *mat.Dense {
	m := mat.NewDense(rows, cols, nil)
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			m.Set(i, j, float64(data[i*cols+j]))
		}
	}
	return m
}

// sliceToGonumVector 将 []float32 切片转换为 gonum 向量
func sliceToGonumVector(data []float32) *mat.VecDense {
	return mat.NewVecDense(len(data), float64s(data))
}

// float32s 将 []float64 转换为 []float32
func float32s(data []float64) []float32 {
	result := make([]float32, len(data))
	for i, v := range data {
		result[i] = float32(v)
	}
	return result
}

// float64s 将 []float32 转换为 []float64
func float64s(data []float32) []float64 {
	result := make([]float64, len(data))
	for i, v := range data {
		result[i] = float64(v)
	}
	return result
}

// randomMatrix 生成随机矩阵
func randomMatrix(rows, cols int) *mat.Dense {
	data := make([]float64, rows*cols)
	for i := range data {
		data[i] = rand.NormFloat64()
	}
	return mat.NewDense(rows, cols, data)
}

// printMatrix 打印矩阵（限制小矩阵）
func printMatrix(name string, m mat.Matrix, rows, cols int) {
	if rows*cols > 100 {
		return // 太大不打印
	}
	fmt.Printf("\n%s (%dx%d):\n", name, rows, cols)
	for i := 0; i < rows; i++ {
		fmt.Print("  ")
		for j := 0; j < cols; j++ {
			fmt.Printf("%8.4f ", m.At(i, j))
		}
		fmt.Println()
	}
}

// printVector 打印向量
func printVector(name string, v []float32, n int) {
	if n > 20 {
		return // 太大不打印
	}
	fmt.Printf("\n%s (len=%d):\n  [", name, n)
	for i := 0; i < n; i++ {
		fmt.Printf("%8.4f", v[i])
		if i < n-1 {
			fmt.Print(", ")
		}
	}
	fmt.Println("]")
}

// computeResidual 计算 ||A * x - b||
func computeResidual(A mat.Matrix, x, b []float32) float64 {
	rows, cols := A.Dims()
	if len(x) != cols || len(b) != rows {
		panic("dimension mismatch in residual")
	}
	maxRes := 0.0
	for i := 0; i < rows; i++ {
		sum := 0.0
		for j := 0; j < cols; j++ {
			sum += A.At(i, j) * float64(x[j])
		}
		res := math.Abs(sum - float64(b[i]))
		if res > maxRes {
			maxRes = res
		}
	}
	return maxRes
}

// testMatrixAddition 测试矩阵加法
func testMatrixAddition(ctx gpucore.GPUContext) bool {
	fmt.Println("\n=== 测试矩阵加法 ===")
	rows, cols := 3, 4
	A := randomMatrix(rows, cols)
	B := randomMatrix(rows, cols)
	expected := mat.NewDense(rows, cols, nil)
	expected.Add(A, B)

	AData := gonumMatrixToSlice(A)
	BData := gonumMatrixToSlice(B)
	A_GPU, err := matrix.NewMatrixFromData(ctx, AData, uint32(rows), uint32(cols))
	if err != nil {
		log.Fatalf("创建矩阵 A 失败: %v", err)
	}
	defer A_GPU.Close()
	B_GPU, err := matrix.NewMatrixFromData(ctx, BData, uint32(rows), uint32(cols))
	if err != nil {
		log.Fatalf("创建矩阵 B 失败: %v", err)
	}
	defer B_GPU.Close()

	C_GPU := A_GPU.Add(B_GPU)
	if err := C_GPU.Err(); err != nil {
		log.Fatalf("矩阵加法失败: %v", err)
	}
	CData, err := C_GPU.Fetch()
	if err != nil {
		log.Fatalf("获取结果失败: %v", err)
	}
	C := sliceToGonumMatrix(CData, rows, cols)

	tol := 1e-5
	isEqual := matrixEqual(C, expected, tol)
	printMatrix("A", A, rows, cols)
	printMatrix("B", B, rows, cols)
	printMatrix("A+B (GPU)", C, rows, cols)
	printMatrix("A+B (Expected)", expected, rows, cols)
	if isEqual {
		fmt.Printf("✅ 矩阵加法测试通过 (容差: %.1e)\n", tol)
	} else {
		fmt.Println("❌ 矩阵加法测试失败")
		maxErr := 0.0
		for i := 0; i < rows; i++ {
			for j := 0; j < cols; j++ {
				err := math.Abs(C.At(i, j) - expected.At(i, j))
				if err > maxErr {
					maxErr = err
				}
			}
		}
		fmt.Printf("最大误差: %.2e\n", maxErr)
	}
	return isEqual
}

// testMatrixMultiplication 测试矩阵乘法
func testMatrixMultiplication(ctx gpucore.GPUContext) bool {
	fmt.Println("\n=== 测试矩阵乘法 ===")
	m, k, n := 3, 4, 2
	A := randomMatrix(m, k)
	B := randomMatrix(k, n)
	expected := mat.NewDense(m, n, nil)
	expected.Mul(A, B)

	AData := gonumMatrixToSlice(A)
	BData := gonumMatrixToSlice(B)
	A_GPU, err := matrix.NewMatrixFromData(ctx, AData, uint32(m), uint32(k))
	if err != nil {
		log.Fatalf("创建矩阵 A 失败: %v", err)
	}
	defer A_GPU.Close()
	B_GPU, err := matrix.NewMatrixFromData(ctx, BData, uint32(k), uint32(n))
	if err != nil {
		log.Fatalf("创建矩阵 B 失败: %v", err)
	}
	defer B_GPU.Close()

	C_GPU := A_GPU.Mul(B_GPU)
	if err := C_GPU.Err(); err != nil {
		log.Fatalf("矩阵乘法失败: %v", err)
	}
	CData, err := C_GPU.Fetch()
	if err != nil {
		log.Fatalf("获取结果失败: %v", err)
	}
	C := sliceToGonumMatrix(CData, m, n)

	tol := 1e-5
	isEqual := matrixEqual(C, expected, tol)
	printMatrix("A", A, m, k)
	printMatrix("B", B, k, n)
	printMatrix("A*B (GPU)", C, m, n)
	printMatrix("A*B (Expected)", expected, m, n)
	if isEqual {
		fmt.Printf("✅ 矩阵乘法测试通过 (容差: %.1e)\n", tol)
	} else {
		fmt.Println("❌ 矩阵乘法测试失败")
		maxErr := 0.0
		for i := 0; i < m; i++ {
			for j := 0; j < n; j++ {
				err := math.Abs(C.At(i, j) - expected.At(i, j))
				if err > maxErr {
					maxErr = err
				}
			}
		}
		fmt.Printf("最大误差: %.2e\n", maxErr)
	}
	return isEqual
}

// testLUFactorization 测试LU分解
func testLUFactorization(ctx gpucore.GPUContext) bool {
	fmt.Println("\n=== 测试LU分解 ===")
	rows := 3

	AData := []float32{
		4, 1, 1,
		1, 3, 2,
		1, 0, 2,
	}

	A := mat.NewDense(rows, rows, float64s(AData))
	var lu mat.LU
	lu.Factorize(A)
	var L, U mat.TriDense
	lu.LTo(&L)
	lu.UTo(&U)

	A_GPU, err := matrix.NewMatrixFromData(ctx, AData, uint32(rows), uint32(rows))
	if err != nil {
		log.Fatalf("创建矩阵 A 失败: %v", err)
	}
	defer A_GPU.Close()

	lu_GPU, pivotBuf, err := A_GPU.LUFactorize()
	if err != nil {
		log.Fatalf("LU分解失败: %v", err)
	}
	defer lu_GPU.Close()
	defer pivotBuf.Release()

	// 【关键】先提交并获取结果，再读取 pivot
	luData, err := lu_GPU.Fetch()
	if err != nil {
		log.Fatalf("获取LU分解结果失败: %v", err)
	}

	// 【关键】在 Fetch 后读取 pivot，此时数据才有效
	if piv, err := matrix.FetchPivot(pivotBuf); err == nil {
		log.Printf("LU分解后 pivot = %v", piv)
	}

	LU := sliceToGonumMatrix(luData, rows, rows)
	printMatrix("A", A, rows, rows)
	printMatrix("L (Expected)", &L, rows, rows)
	printMatrix("U (Expected)", &U, rows, rows)
	printMatrix("LU (GPU)", LU, rows, rows)

	L_gpu := mat.NewTriDense(rows, false, nil)
	U_gpu := mat.NewTriDense(rows, true, nil)
	for i := 0; i < rows; i++ {
		for j := 0; j < rows; j++ {
			if i > j {
				L_gpu.SetTri(i, j, float64(LU.At(i, j)))
			} else {
				U_gpu.SetTri(i, j, float64(LU.At(i, j)))
			}
		}
		L_gpu.SetTri(i, i, 1.0)
	}

	var A_reconstructed_gpu mat.Dense
	A_reconstructed_gpu.Mul(L_gpu, U_gpu)

	tol := 1e-4
	isEqual := matrixEqual(&A_reconstructed_gpu, A, tol)
	if isEqual {
		fmt.Printf("✅ LU分解测试通过 (容差: %.1e)\n", tol)
	} else {
		fmt.Printf("❌ LU分解测试失败\n")
	}
	return isEqual
}

// testLinearSolve 测试线性方程求解
// testLinearSolve 测试线性方程求解
func testLinearSolve(ctx gpucore.GPUContext) bool {
	fmt.Println("\n=== 测试线性方程求解 ===")
	allPassed := true

	// === 测试用例1：简单对角系统（理论解明确）===
	fmt.Println("\n测试用例1：简单对角系统")
	{
		rows := 3
		AData := []float32{
			4, 0, 0,
			0, 3, 0,
			0, 0, 2,
		}
		bData := []float32{8, 6, 4} // 解应为 [2, 2, 2]

		A_GPU, err := matrix.NewMatrixFromData(ctx, AData, uint32(rows), uint32(rows))
		if err != nil {
			log.Printf("❌ 创建矩阵失败: %v", err)
			allPassed = false
			goto case2
		}
		defer A_GPU.Close()

		b_GPU, err := matrix.NewMatrixFromData(ctx, bData, uint32(rows), 1)
		if err != nil {
			log.Printf("❌ 创建向量失败: %v", err)
			allPassed = false
			goto case2
		}
		defer b_GPU.Close()

		// LU 分解（不提交，保持 encoder 活跃）
		lu_GPU, pivotBuf, err := A_GPU.LUFactorize()
		if err != nil {
			log.Printf("❌ LU分解失败: %v", err)
			allPassed = false
			goto case2
		}
		defer lu_GPU.Close()
		defer pivotBuf.Release()

		// 求解（复用同一个 encoder）
		x_GPU := lu_GPU.Solve(b_GPU)
		xData, err := x_GPU.Fetch() // Submit 在这里发生
		if err != nil {
			log.Printf("❌ 求解失败: %v", err)
			allPassed = false
			goto case2
		}

		// 安全读取 pivot（必须在 Fetch 后，GPU 操作已完成）
		piv, _ := matrix.FetchPivot(pivotBuf)
		log.Printf("测试用例1: pivot=%v, x=%v", piv, xData)

		// 验证残差
		A := mat.NewDense(rows, rows, float64s(AData))
		maxResidual := computeResidual(A, xData, bData)
		fmt.Printf("最大残差: %.2e\n", maxResidual)

		if maxResidual > 1e-3 {
			fmt.Printf("❌ 测试用例1失败（残差过高）\n")
			allPassed = false
		} else {
			fmt.Printf("✅ 测试用例1通过\n")
		}
	}

case2:
	// === 测试用例2：原始测试矩阵 ===
	fmt.Println("\n测试用例2：原始测试矩阵")
	{
		rows := 3
		AData := []float32{
			2, 1, 1,
			1, 3, 2,
			1, 0, 0,
		}
		bData := []float32{4, 5, 1} // 解应为 [1, 0, 2]

		A_GPU, err := matrix.NewMatrixFromData(ctx, AData, uint32(rows), uint32(rows))
		if err != nil {
			log.Printf("❌ 创建矩阵失败: %v", err)
			allPassed = false
			goto case3
		}
		defer A_GPU.Close()

		b_GPU, err := matrix.NewMatrixFromData(ctx, bData, uint32(rows), 1)
		if err != nil {
			log.Printf("❌ 创建向量失败: %v", err)
			allPassed = false
			goto case3
		}
		defer b_GPU.Close()

		lu_GPU, pivotBuf, err := A_GPU.LUFactorize()
		if err != nil {
			log.Printf("❌ LU分解失败: %v", err)
			allPassed = false
			goto case3
		}
		defer lu_GPU.Close()
		defer pivotBuf.Release()

		x_GPU := lu_GPU.Solve(b_GPU)
		xData, err := x_GPU.Fetch()
		if err != nil {
			log.Printf("❌ 求解失败: %v", err)
			allPassed = false
			goto case3
		}

		piv, _ := matrix.FetchPivot(pivotBuf)
		log.Printf("测试用例2: pivot=%v, x=%v", piv, xData)

		A := mat.NewDense(rows, rows, float64s(AData))
		maxResidual := computeResidual(A, xData, bData)
		fmt.Printf("最大残差: %.2e\n", maxResidual)

		if maxResidual > 1e-3 {
			fmt.Printf("❌ 测试用例2失败（残差过高）\n")
			allPassed = false
		} else {
			fmt.Printf("✅ 测试用例2通过\n")
		}
	}

case3:
	// === 测试用例3：随机矩阵（对角占优）===
	fmt.Println("\n测试用例3：随机矩阵")
	{
		n := 5

		// 生成对角占优矩阵确保可解
		AData := make([]float32, n*n)
		for i := 0; i < n; i++ {
			for j := 0; j < n; j++ {
				if i == j {
					AData[i*n+j] = float32(n) + rand.Float32()*2
				} else {
					AData[i*n+j] = (rand.Float32() - 0.5) / float32(n)
				}
			}
		}
		bData := make([]float32, n)
		for i := range bData {
			bData[i] = rand.Float32() * 10
		}

		A_GPU, err := matrix.NewMatrixFromData(ctx, AData, uint32(n), uint32(n))
		if err != nil {
			log.Printf("❌ 创建矩阵失败: %v", err)
			allPassed = false
			goto case4
		}
		defer A_GPU.Close()

		b_GPU, err := matrix.NewMatrixFromData(ctx, bData, uint32(n), 1)
		if err != nil {
			log.Printf("❌ 创建向量失败: %v", err)
			allPassed = false
			goto case4
		}
		defer b_GPU.Close()

		lu_GPU, pivotBuf, err := A_GPU.LUFactorize()
		if err != nil {
			log.Printf("❌ LU分解失败: %v", err)
			allPassed = false
			goto case4
		}
		defer lu_GPU.Close()
		defer pivotBuf.Release()

		x_GPU := lu_GPU.Solve(b_GPU)
		xData, err := x_GPU.Fetch()
		if err != nil {
			log.Printf("❌ 求解失败: %v", err)
			allPassed = false
			goto case4
		}

		piv, _ := matrix.FetchPivot(pivotBuf)
		log.Printf("测试用例3: pivot=%v, x[:5]=%v", piv, xData[:min(5, len(xData))])

		A := mat.NewDense(n, n, float64s(AData))
		maxResidual := computeResidual(A, xData, bData)
		fmt.Printf("最大残差: %.2e\n", maxResidual)

		if maxResidual > 1e-2 {
			fmt.Printf("❌ 测试用例3失败（残差过高）\n")
			allPassed = false
		} else {
			fmt.Printf("✅ 测试用例3通过\n")
		}
	}

case4:
	// === 测试用例4：奇异矩阵检测 ===
	fmt.Println("\n测试用例4：奇异矩阵")
	{
		singularData := []float32{
			1, 2, 3,
			4, 5, 6,
			7, 8, 9,
		}
		// singularB := []float32{1, 2, 3}

		singularA_GPU, err := matrix.NewMatrixFromData(ctx, singularData, 3, 3)
		if err != nil {
			log.Printf("❌ 创建奇异矩阵失败: %v", err)
			allPassed = false
			goto end
		}
		defer singularA_GPU.Close()

		lu_singular, pivotBuf_singular, err := singularA_GPU.LUFactorize()
		if err != nil {
			log.Printf("❌ 奇异矩阵LU分解失败: %v", err)
			allPassed = false
			goto end
		}
		defer lu_singular.Close()
		defer pivotBuf_singular.Release()

		// 【关键】使用新的方法检查奇异性
		log.Printf("正在检查矩阵奇异性...")
		if err := lu_singular.CheckSingularityAfterLU(); err != nil {
			if err.Error() == "matrix is singular" { // 修复：匹配实际返回的错误字符串
				fmt.Printf("✅ 奇异矩阵正确检测: %v\n", err)
			} else {
				log.Printf("❌ 奇异矩阵检测错误: %v", err)
				allPassed = false
			}
		} else {
			fmt.Printf("❌ 奇异矩阵未被检测\n")
			allPassed = false
		}

		// 【新增】验证 pivot 缓冲区的内容
		if piv, err := matrix.FetchPivot(pivotBuf_singular); err != nil {
			log.Printf("✅ FetchPivot 也正确报告了错误: %v\n", err)
		} else {
			log.Printf("❌ FetchPivot 应该失败但返回了: %v\n", piv)
			allPassed = false
		}
	}

end:
	return allPassed
}

// 辅助函数：取最小值
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// testLargeMatrix 测试大矩阵运算
func testLargeMatrix(ctx gpucore.GPUContext) bool {
	fmt.Println("\n=== 测试大矩阵运算 ===")
	size := 500 // 【关键】增大到 500x500 以体现 GPU 优势
	fmt.Printf("创建 %dx%d 矩阵...\n", size, size)

	A := randomMatrix(size, size)
	B := randomMatrix(size, size)

	startCPU := time.Now()
	var C_expected mat.Dense
	C_expected.Mul(A, B)
	cpuTime := time.Since(startCPU)

	AData := gonumMatrixToSlice(A)
	BData := gonumMatrixToSlice(B)

	startUpload := time.Now()
	A_GPU, err := matrix.NewMatrixFromData(ctx, AData, uint32(size), uint32(size))
	if err != nil {
		log.Fatalf("创建大矩阵 A 失败: %v", err)
	}
	defer A_GPU.Close()
	B_GPU, err := matrix.NewMatrixFromData(ctx, BData, uint32(size), uint32(size))
	if err != nil {
		log.Fatalf("创建大矩阵 B 失败: %v", err)
	}
	defer B_GPU.Close()
	uploadTime := time.Since(startUpload)

	startGPU := time.Now()
	C_GPU := A_GPU.Mul(B_GPU)
	if err := C_GPU.Err(); err != nil {
		log.Fatalf("大矩阵乘法失败: %v", err)
	}
	CData, err := C_GPU.Fetch()
	if err != nil {
		log.Fatalf("获取大矩阵结果失败: %v", err)
	}
	gpuTime := time.Since(startGPU)

	C := sliceToGonumMatrix(CData, size, size)

	// 对于大矩阵，适当放宽容差
	tol := 1e-2
	isEqual := matrixEqual(C, &C_expected, tol)
	if isEqual {
		fmt.Printf("✅ %dx%d 大矩阵乘法测试通过 (容差: %.1e)\n", size, size, tol)
	} else {
		fmt.Printf("❌ %dx%d 大矩阵乘法测试失败\n", size, size)
	}

	fmt.Printf("CPU 计算时间: %v\n", cpuTime)
	fmt.Printf("GPU 上传时间: %v\n", uploadTime)
	fmt.Printf("GPU 计算时间: %v\n", gpuTime)
	speedup := float64(cpuTime) / float64(gpuTime)
	if speedup > 1.0 {
		fmt.Printf("GPU 比 CPU 快: %.2fx\n", speedup)
	} else {
		fmt.Printf("CPU 比 GPU 快: %.2fx\n", 1.0/speedup)
	}
	return isEqual
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	rand.Seed(time.Now().UnixNano())

	ctx, err := gpucore.NewGPUContext()
	if err != nil {
		log.Fatalf("Failed to create GPU context: %v", err)
	}
	defer ctx.Close()

	if err := compute.RegisterKernels(ctx); err != nil {
		log.Fatalf("Failed to register kernels: %v", err)
	}

	fmt.Println("开始矩阵运算测试...")
	allTestsPassed := true

	if !testMatrixAddition(ctx) {
		allTestsPassed = false
	}
	if !testMatrixMultiplication(ctx) {
		allTestsPassed = false
	}
	if !testLUFactorization(ctx) {
		allTestsPassed = false
	}
	if !testLinearSolve(ctx) {
		allTestsPassed = false
	}
	if !testLargeMatrix(ctx) {
		allTestsPassed = false
	}

	fmt.Println("\n====================")
	if allTestsPassed {
		fmt.Println("🎉 所有测试通过！")
		os.Exit(0)
	} else {
		fmt.Println("❌ 有测试失败！")
		os.Exit(1)
	}
}
